package com.alibaba.datax.plugin.reader.elasticsearchreader;

import com.alibaba.datax.common.element.*;
import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.common.plugin.RecordSender;
import com.alibaba.datax.common.spi.Reader;
import com.alibaba.datax.common.util.Configuration;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.action.search.*;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class EsReader extends Reader {
    public static class Job extends Reader.Job {
        private static final Logger LOG = LoggerFactory.getLogger(Job.class);
        private Configuration originalConfig = null;

        @Override
        public void init() {
            this.originalConfig = super.getPluginJobConf();
            validateConfig();
        }

        private void validateConfig() {
            this.originalConfig.getNecessaryValue(Key.ENDPOINT, EsReaderErrorCode.REQUIRED_VALUE);
            this.originalConfig.getNecessaryValue(Key.INDEX, EsReaderErrorCode.REQUIRED_VALUE);

            if (Key.isIncremental(this.originalConfig)) {
                this.originalConfig.getNecessaryValue(Key.INCREMENTAL_COLUMN, EsReaderErrorCode.REQUIRED_VALUE);
            }
        }

        @Override
        public void prepare() {
        }

        @Override
        public List<Configuration> split(int adviceNumber) {
            List<Configuration> configurations = new ArrayList<>();

            // 如果是增量模式，根据增量字段进行范围划分
            if (Key.isIncremental(this.originalConfig)) {
                // TODO: 实现增量拆分逻辑
                configurations.add(this.originalConfig);
            } else {
                // 全量模式直接返回原配置
                configurations.add(this.originalConfig);
            }

            return configurations;
        }

        @Override
        public void post() {
        }

        @Override
        public void destroy() {
        }
    }

    public static class Task extends Reader.Task {
        private static final Logger LOG = LoggerFactory.getLogger(Task.class);
        private Configuration readerConfig;
        private RestHighLevelClient client;
        private String scrollId;
        private List<String> columnList;

        @Override
        public void init() {
            this.readerConfig = super.getPluginJobConf();
            this.columnList = this.readerConfig.getList(Key.COLUMN, String.class);
            initClient();
        }

        private void initClient() {
            try {
                URI uri = new URI(Key.getEndpoint(readerConfig));
                HttpHost httpHost = new HttpHost(uri.getHost(), uri.getPort(), uri.getScheme());

                RestClientBuilder builder = RestClient.builder(httpHost);

                // 设置认证信息
                String username = Key.getUsername(readerConfig);
                String password = Key.getPassword(readerConfig);
                if (username != null && password != null) {
                    final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
                    credentialsProvider.setCredentials(AuthScope.ANY,
                            new UsernamePasswordCredentials(username, password));
                    builder.setHttpClientConfigCallback(httpClientBuilder ->
                            httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
                }

                this.client = new RestHighLevelClient(builder);
            } catch (URISyntaxException e) {
                throw DataXException.asDataXException(EsReaderErrorCode.CONFIG_INVALID_EXCEPTION,
                        "Invalid endpoint: " + Key.getEndpoint(readerConfig), e);
            }
        }

        @Override
        public void prepare() {
        }

        @Override
        public void startRead(RecordSender recordSender) {
            try {
                // 创建scroll上下文
                final Scroll scroll = new Scroll(TimeValue.timeValueMinutes(1));
                SearchRequest searchRequest = new SearchRequest(Key.getIndex(readerConfig));
                searchRequest.scroll(scroll);

                // 构建查询
                SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
                if (Key.isIncremental(readerConfig)) {
                    // 增量查询
                    String incrementalColumn = Key.getIncrementalColumn(readerConfig);
                    String lastValue = Key.getLastValue(readerConfig);
                    searchSourceBuilder.query(
                            QueryBuilders.rangeQuery(incrementalColumn).gt(lastValue)
                    );
                } else {
                    // 全量查询
                    String searchQuery = Key.getSearchQuery(readerConfig);
                    try {
                        if (searchQuery != null && !searchQuery.trim().isEmpty()) {
                            JSONObject queryObject = JSON.parseObject(searchQuery);
                            if (queryObject != null && queryObject.containsKey("query")) {
                                JSONObject query = queryObject.getJSONObject("query");
                                // 根据查询类型构建对应的QueryBuilder
                                if (query.containsKey("match_all")) {
                                    searchSourceBuilder.query(QueryBuilders.matchAllQuery());
                                } else if (query.containsKey("match")) {
                                    JSONObject match = query.getJSONObject("match");
                                    String field = match.keySet().iterator().next();
                                    searchSourceBuilder.query(QueryBuilders.matchQuery(field, match.get(field)));
                                } else if (query.containsKey("term")) {
                                    JSONObject term = query.getJSONObject("term");
                                    String field = term.keySet().iterator().next();
                                    searchSourceBuilder.query(QueryBuilders.termQuery(field, term.get(field)));
                                } else if (query.containsKey("range")) {
                                    JSONObject range = query.getJSONObject("range");
                                    String field = range.keySet().iterator().next();
                                    JSONObject rangeParams = range.getJSONObject(field);
                                    RangeQueryBuilder rangeQuery = QueryBuilders.rangeQuery(field);
                                    if (rangeParams.containsKey("gt")) {
                                        rangeQuery.gt(rangeParams.get("gt"));
                                    }
                                    if (rangeParams.containsKey("gte")) {
                                        rangeQuery.gte(rangeParams.get("gte"));
                                    }
                                    if (rangeParams.containsKey("lt")) {
                                        rangeQuery.lt(rangeParams.get("lt"));
                                    }
                                    if (rangeParams.containsKey("lte")) {
                                        rangeQuery.lte(rangeParams.get("lte"));
                                    }
                                    searchSourceBuilder.query(rangeQuery);
                                } else {
                                    throw DataXException.asDataXException(EsReaderErrorCode.CONFIG_INVALID_EXCEPTION,
                                            "Unsupported query type. Supported types are: match_all, match, term, range");
                                }
                            } else {
                                throw DataXException.asDataXException(EsReaderErrorCode.CONFIG_INVALID_EXCEPTION,
                                        "Invalid search query format: missing 'query' field");
                            }
                        } else {
                            // 默认使用match_all查询
                            searchSourceBuilder.query(QueryBuilders.matchAllQuery());
                        }
                    } catch (Exception e) {
                        throw DataXException.asDataXException(EsReaderErrorCode.CONFIG_INVALID_EXCEPTION,
                                "Failed to parse search query: " + e.getMessage(), e);
                    }
                }

                searchSourceBuilder.size(Key.getScrollSize(readerConfig));
                searchRequest.source(searchSourceBuilder);

                // 执行首次查询
                SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
                scrollId = searchResponse.getScrollId();
                processSearchResponse(searchResponse, recordSender);

                // 循环处理后续批次
                while (true) {
                    SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
                    scrollRequest.scroll(scroll);
                    searchResponse = client.scroll(scrollRequest, RequestOptions.DEFAULT);

                    if (searchResponse.getHits().getHits().length == 0) {
                        break;
                    }

                    processSearchResponse(searchResponse, recordSender);
                }
            } catch (IOException e) {
                throw DataXException.asDataXException(EsReaderErrorCode.ES_SEARCH_ERROR, e);
            } finally {
                // 清理scroll上下文
                if (scrollId != null) {
                    try {
                        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
                        clearScrollRequest.addScrollId(scrollId);
                        client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
                    } catch (IOException e) {
                        LOG.warn("Failed to clear scroll context", e);
                    }
                }
            }
        }

        private void processSearchResponse(SearchResponse searchResponse, RecordSender recordSender) {
            for (SearchHit hit : searchResponse.getHits().getHits()) {
                Record record = recordSender.createRecord();
                Map<String, Object> source = hit.getSourceAsMap();

                if (columnList == null || columnList.isEmpty()) {
                    // 如果没有指定列，则获取所有列
                    for (Map.Entry<String, Object> entry : source.entrySet()) {
                        record.addColumn(convertToColumn(entry.getValue()));
                    }
                } else {
                    // 按指定列获取数据
                    for (String columnName : columnList) {
                        Object value = source.get(columnName);
                        record.addColumn(convertToColumn(value));
                    }
                }

                recordSender.sendToWriter(record);
            }
        }

        private Column convertToColumn(Object value) {
            if (value == null) {
                return new StringColumn(null);
            }

            if (value instanceof String) {
                return new StringColumn((String) value);
            } else if (value instanceof Integer) {
                return new LongColumn((Integer) value);
            } else if (value instanceof Long) {
                return new LongColumn((Long) value);
            } else if (value instanceof Float) {
                return new DoubleColumn((Float) value);
            } else if (value instanceof Double) {
                return new DoubleColumn((Double) value);
            } else if (value instanceof Boolean) {
                return new BoolColumn((Boolean) value);
            } else if (value instanceof Map || value instanceof List) {
                return new StringColumn(JSON.toJSONString(value));
            } else {
                return new StringColumn(value.toString());
            }
        }

        @Override
        public void post() {
        }

        @Override
        public void destroy() {
            if (client != null) {
                try {
                    client.close();
                } catch (IOException e) {
                    LOG.warn("Failed to close ES client", e);
                }
            }
        }
    }
}
